from sklearn.cluster import KMeans
import cv2
import matplotlib.pyplot as plt

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  
plt.rcParams['axes.unicode_minus'] = False  

# 读取图像
image = cv2.imread('data/01.jpeg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 将图像转换为二维数组
pixels = image.reshape(-1, 3)

# 使用K-means进行颜色聚类
kmeans = KMeans(n_clusters=16, random_state=42)
labels = kmeans.fit_predict(pixels)

# 用聚类中心替换原像素
compressed_image = kmeans.cluster_centers_[labels].reshape(image.shape)
compressed_image = compressed_image.astype('uint8')

# 显示结果
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.imshow(image)
plt.title('原始图像')
plt.axis('off')

plt.subplot(1, 2, 2)
plt.imshow(compressed_image)
plt.title('16色量化图像')
plt.axis('off')
plt.show()